import argparse
import sys

def classify_loops(input_bedpe, output_1mb, output_1_2mb, output_2mb):
    """分类处理BEDPE文件中的loop距离"""
    with open(input_bedpe, 'r') as infile, \
         open(output_1mb, 'w') as file_1mb, \
         open(output_1_2mb, 'w') as file_1_2mb, \
         open(output_2mb, 'w') as file_2mb:
        
        # 处理表头（保留所有注释行和表头行）
        header_lines = []
        while True:
            pos = infile.tell()
            line = infile.readline()
            
            # 文件结束或遇到数据行
            if not line or not line.startswith('#'):
                infile.seek(pos)  # 回退到数据行开始位置
                print(1)
                break
                
            header_lines.append(line)
        
        # 将表头写入所有输出文件
        for hline in header_lines:
            file_1mb.write(hline)
            file_1_2mb.write(hline)
            file_2mb.write(hline)
        
        # 处理数据行
        for line in infile:
            parts = line.strip().split()
            print('len: ', len(parts))
            if len(parts) < 24:  # 确保有足够的列包含centroid信息
                file_2mb.write(line)
                continue
                
            chrom1, chrom2 = parts[0], parts[3]
            
            # 尝试获取centroid位置（倒数第3和倒数第2列）
            try:
                centroid1 = int(float(parts[-3]))  # 处理可能的科学计数法
                centroid2 = int(float(parts[-2]))
            except (ValueError, IndexError):
                # 如果centroid列无效，尝试使用锚点中心位置
                try:
                    start1, end1 = int(parts[1]), int(parts[2])
                    start2, end2 = int(parts[4]), int(parts[5])
                    centroid1 = (start1 + end1) // 2
                    centroid2 = (start2 + end2) // 2
                except (ValueError, IndexError):
                    # 如果仍然失败，归类为长距离loop
                    file_2mb.write(line)
                    continue
            
            # 跨染色体或距离计算
            if chrom1 != chrom2:
                file_2mb.write(line)
            else:
                distance = abs(centroid2 - centroid1)
                print(distance)
                if distance <= 1000000:  # ≤1Mb
                    file_1mb.write(line)
                elif distance <= 2000000:  # 1Mb-2Mb
                    file_1_2mb.write(line)
                else:  # >2Mb
                    file_2mb.write(line)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Classify chromatin loops by distance')
    parser.add_argument('-i', '--input', required=True, help='Input BEDPE file')
    parser.add_argument('-a', '--within_1mb', required=True, help='Output for loops ≤1Mb')
    parser.add_argument('-b', '--between_1_2mb', required=True, help='Output for loops 1Mb-2Mb')
    parser.add_argument('-c', '--beyond_2mb', required=True, help='Output for loops >2Mb or inter-chromosomal')
    
    args = parser.parse_args()
    
    classify_loops(args.input, args.within_1mb, args.between_1_2mb, args.beyond_2mb)
    
    print("Loop classification completed:")
    print(f"- ≤1Mb loops saved to: {args.within_1mb}")
    print(f"- 1Mb-2Mb loops saved to: {args.between_1_2mb}")
    print(f"- >2Mb/inter-chromosomal loops saved to: {args.beyond_2mb}")